import torch
from torch import nn

class PositionWiseFFN(nn.Module):
    def __init__(self, input_dim):
        super(PositionWiseFFN, self).__init__()
        # 强令隐藏层维度是输入的4倍
        self.Linear1=nn.Linear(input_dim, input_dim*4)
        self.Linear2=nn.Linear(input_dim*4, input_dim)
        self.activate=nn.ReLU()
        
    def forward(self, x):
        return self.Linear2(self.activate(self.Linear1(x)))
    
class PositionalEncoding(nn.Module):
    """Positional encoding.

    Defined in :numref:`sec_self-attention-and-positional-encoding`"""
    def __init__(self, num_hiddens, dropout=0, max_len=1000):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        # Create a long enough P
        self.P = torch.zeros([1, max_len, num_hiddens])
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)

if __name__=='__main__':
    x=torch.zeros([2,10,512])
    my_encoding=PositionalEncoding(num_hiddens=512, dropout=0)
    my_encoding(x)
    #print(my_encoding(x))
    